#include "stdafx.h"
#include "MemoryArea.h"
#include <algorithm>

namespace MiniBasic
{
    MemoryArea::MemoryArea(const SIZE_T& bytesToAllocate)
    {
        currentPointer = base = reinterpret_cast<BYTE*>(VirtualAlloc(NULL, std::max<SIZE_T>(bytesToAllocate, 1), MEM_COMMIT, PAGE_READWRITE));
        if(!base)
        {
            throw E_OUTOFMEMORY;
        }
        end = base + bytesToAllocate;
    }
    
    MemoryArea::~MemoryArea()
    {
        VirtualFree(base, 0, MEM_FREE);
    }

    ULONG_PTR MemoryArea::Size() const
    {
        return end - base;
    }

    STDMETHODIMP MemoryArea::QueryInterface(REFIID iid, void** ppvObject)
    { 
        if((iid == IID_IUnknown) || (iid == IID_IStream) || (iid == IID_ISequentialStream))
        {
            *ppvObject = static_cast<IStream*>(this);
            AddRef();
            return S_OK;
        } 
        else
        {
            *ppvObject = NULL;
            return E_NOINTERFACE;
        }
    }

    STDMETHODIMP_(ULONG) MemoryArea::AddRef(void) 
    { 
        return (ULONG)InterlockedIncrement(&refCount); 
    }

    STDMETHODIMP_(ULONG) MemoryArea::Release(void) 
    {
        ULONG res = (ULONG) InterlockedDecrement(&refCount);
        if(res == 0)
        {
            delete this;
        }
        return res;
    }

    // ISequentialStream Interface
    STDMETHODIMP MemoryArea::Read(void* pv, ULONG cb, ULONG* pcbRead)
    {
        HRESULT hr = S_OK;
        if(pv)
        {
            // make sure we don't go past the end
            ULONG bytesToCopy = std::min<ULONG>(end - currentPointer, cb);
            if(bytesToCopy)
            {
                std::copy(currentPointer, currentPointer + bytesToCopy, static_cast<BYTE*>(pv));
                currentPointer += bytesToCopy;
                if(pcbRead)
                {
                    *pcbRead = bytesToCopy;
                }
            }
            else
            {
                hr = STG_E_MEDIUMFULL;
            }
        }
        else
        {
            hr = STG_E_INVALIDPOINTER;
        }
        return S_OK;
    }

    STDMETHODIMP MemoryArea::Write(void const* pv, ULONG cb, ULONG* pcbWritten)
    {
        HRESULT hr = S_OK;
        if(pv)
        {
            // make sure we don't go past the end
            ULONG bytesToCopy = std::min<ULONG>(end - currentPointer, cb);
            if(bytesToCopy)
            {
                const BYTE* input = static_cast<const BYTE*>(pv);
                std::copy(input, input + bytesToCopy, static_cast<BYTE*>(currentPointer));
                currentPointer += bytesToCopy;
                if(pcbWritten)
                {
                    *pcbWritten = bytesToCopy;
                }
            }
            else
            {
                hr = STG_E_MEDIUMFULL;
            }
        }
        else
        {
            hr = STG_E_INVALIDPOINTER;
        }
        return hr;
    }

    STDMETHODIMP MemoryArea::SetSize(ULARGE_INTEGER ul)
    {
        HRESULT hr = S_OK;
        if(ul.QuadPart <= Size())
        {
            end = base + ul.QuadPart;
            if(currentPointer > end)
            {
                currentPointer = end;
            }
        }
        else
        {
            BYTE* pNewMem = static_cast<BYTE*>(VirtualAlloc(NULL, static_cast<SIZE_T>(ul.QuadPart), MEM_COMMIT, PAGE_READWRITE));
            if(pNewMem)
            {
                memcpy(pNewMem, base, Size());
                BYTE* currentBaseCopy = base;
                BYTE* newEnd = pNewMem + ul.QuadPart;
                BYTE* newCurrent = pNewMem + (currentPointer - base);
                base = pNewMem;
                end = newEnd;
                currentPointer = newCurrent;
                VirtualFree(currentBaseCopy, 0, MEM_FREE);
            }
            else
            {
                hr = E_OUTOFMEMORY;
            }
        }
        return hr;   
    }
    
    STDMETHODIMP MemoryArea::CopyTo(IStream* pstrm, ULARGE_INTEGER cb, ULARGE_INTEGER* pcbRead, ULARGE_INTEGER* pcbWritten)
    { 
        HRESULT hr = S_OK;
        if(pstrm)
        {
            ULONGLONG totalBytesToCopy = std::min<ULONGLONG>(cb.QuadPart, end - currentPointer);
            ULONGLONG bytesCopied = 0;
            while(totalBytesToCopy)
            {
                ULONG bytesToCopy = static_cast<ULONG>(std::min<ULONGLONG>(MAXDWORD, totalBytesToCopy));
                ULONG copiedThisOperation = 0;
                hr = pstrm->Write(currentPointer + bytesCopied, bytesToCopy, &copiedThisOperation);
                if(SUCCEEDED(hr) || (copiedThisOperation != 0))
                {
                    bytesCopied += copiedThisOperation;
                    totalBytesToCopy -= copiedThisOperation;
                }
                else break;
            }
            if(pcbRead)
            {
                pcbRead->QuadPart = bytesCopied;
            }
            if(pcbWritten)
            {
                pcbWritten->QuadPart = bytesCopied;
            }
            if(cb.QuadPart != MAXULONGLONG)
            {
                hr = totalBytesToCopy == bytesCopied ? S_OK : hr;
            }
        }
        else
        {
            hr = STG_E_INVALIDPOINTER;
        }
        return hr;   
    }

    STDMETHODIMP MemoryArea::Commit(DWORD)                                      
    { 
        return E_NOTIMPL;   
    }
    
    STDMETHODIMP MemoryArea::Revert(void)                                       
    { 
        return E_NOTIMPL;   
    }
    
    STDMETHODIMP MemoryArea::LockRegion(ULARGE_INTEGER, ULARGE_INTEGER, DWORD)              
    { 
        return E_NOTIMPL;   
    }
    
    STDMETHODIMP MemoryArea::UnlockRegion(ULARGE_INTEGER, ULARGE_INTEGER, DWORD)            
    { 
        return E_NOTIMPL;   
    }
    
    STDMETHODIMP MemoryArea::Clone(IStream **)                                  
    { 
        return E_NOTIMPL;   
    }

    STDMETHODIMP MemoryArea::Seek(LARGE_INTEGER liDistanceToMove, DWORD dwOrigin, ULARGE_INTEGER* lpNewFilePointer)
    { 
        BYTE* seekPointer = NULL;
        bool asUnsigned = false;
        ULARGE_INTEGER uli = {0};

        switch(dwOrigin)
        {
            case STREAM_SEEK_SET:
            {
                asUnsigned = true;
                uli.QuadPart = static_cast<ULONGLONG>(liDistanceToMove.QuadPart);
                seekPointer = base;
            }
            break;
            case STREAM_SEEK_CUR:
            {
                seekPointer = currentPointer;
            }
            break;
            case STREAM_SEEK_END:
            {
                seekPointer = end;
            }
            break;
            default:
            {
                return STG_E_INVALIDFUNCTION;
            }
            break;
        }

        BYTE* temp = seekPointer + (asUnsigned ? uli.QuadPart : liDistanceToMove.QuadPart);
        if(temp >= base && temp <= end)
        {
            currentPointer = temp;
            if(lpNewFilePointer)
            {
                lpNewFilePointer->QuadPart = static_cast<ULONGLONG>(currentPointer - base);
            }
            return S_OK;
        }
        else
        {
            return HRESULT_FROM_WIN32(ERROR_SEEK_ON_DEVICE);
        }
    }

    STDMETHODIMP MemoryArea::Stat(STATSTG* pStatstg, DWORD) 
    {
        HRESULT hr = S_OK;
        if(pStatstg)
        {
            std::fill_n(reinterpret_cast<char*>(pStatstg), sizeof(STATSTG), 0);
            pStatstg->cbSize.QuadPart = Size();
            pStatstg->type = STGTY_LOCKBYTES;
        }
        else
        {
            hr = STG_E_INVALIDPOINTER;
        }
        return S_OK;
    }
}
